from collections import deque

import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from tqdm import tqdm

import config
import modeldef
import preprocess


class TrainItem:
    def __init__(self, now_epoch: int, train_loss: float, train_acc: float, val_acc: float):
        self.now_epoch = now_epoch
        self.train_loss = train_loss
        self.train_acc = train_acc
        self.val_acc = val_acc


train_queues = {}
total_epochs = {}


def stage_test(dataloader, model, device):
    corrects = 0
    epoch_sum = 0
    for _, (x, y) in enumerate(dataloader):
        x = x.to(device)
        y = y.to(device)
        model.eval()
        output = model(x)
        label = torch.argmax(output.cpu(), 1)
        epoch_sum += x.shape[0]
        corrects += torch.sum(label == y.cpu().data)
    return corrects / epoch_sum


def get_optimizer(optimizer_name: str, model: nn.Module, *args, **kwargs) -> torch.optim.Optimizer:
    if optimizer_name == "adam":
        return torch.optim.Adam(model.parameters(), *args, **kwargs)
    elif optimizer_name == "sgd":
        return torch.optim.SGD(model.parameters(), *args, **kwargs)
    else:
        raise ValueError("Unknown optimizer type")


def train(model_type="lenet", model_name="lenet", dataset=1, lr=0.01, num_epochs=50, batch_size=32,
          optimizer="adam", device="cuda"):
    global train_queues
    if device == "cuda" and not torch.cuda.is_available():
        device = "cpu"
    print(f"Training on {device}.")
    model = modeldef.get_model(model_type).to(device)
    optim = get_optimizer(optimizer, model, lr)
    loss_func = nn.CrossEntropyLoss()
    train_loader, valid_loader, test_loader = preprocess.get_dataloader(dataset, batch_size=batch_size)
    print(
        f"train data number: {len(train_loader)}, valid data number: {len(valid_loader)}, test data number: {len(test_loader)}")
    epoch_loss = 0
    epoch_corrects = 0
    epoch_sum = 0
    train_queues[model_name] = deque()
    total_epochs[model_name] = num_epochs
    for epoch in range(num_epochs):
        print("Epoch {}/{}".format(epoch, num_epochs))
        for i, (x, y) in tqdm(enumerate(train_loader)):
            # print(f"\titer {i}")
            x = x.to(device)
            y = y.to(device)
            model.train()
            output = model(x)
            loss = loss_func(output, y)
            optim.zero_grad()
            loss.backward()
            optim.step()

            label = torch.argmax(output.cpu(), 1)
            epoch_loss += loss.item() * x.shape[0]
            epoch_corrects += torch.sum(label == y.cpu())
            epoch_sum += x.shape[0]

        cur_train_loss = epoch_loss / epoch_sum
        cur_train_acc = epoch_corrects / epoch_sum

        cur_valid_acc = stage_test(valid_loader, model, device)

        print(f"train loss: {cur_train_loss}, train_acc: {cur_train_acc}, valid_acc: {cur_valid_acc}")
        train_queues[model_name].append(
            TrainItem(epoch + 1, cur_train_loss, cur_train_acc.item(), cur_valid_acc.item()))
    test_acc = stage_test(test_loader, model, device)
    print(f"test_acc: {test_acc}")

    torch.save(model, config.get_model_path(model_name))
    return test_acc


# 预处理函数

def preProccessing(img):
    imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    imgBlur = cv2.GaussianBlur(imgGray, (5, 5), 1)
    imgCanny = cv2.Canny(imgBlur, 200, 200)
    imgDial = cv2.dilate(imgCanny, np.ones((5, 5)), iterations=2)  # 膨胀操作
    imgThres = cv2.erode(imgDial, np.ones((5, 5)), iterations=1)  # 腐蚀操作
    return imgThres


def getContours(img, input_img):
    x, y, w, h, xx, yy, ss = 0, 0, 10, 10, 20, 20, 10  # 因为图像大小不能为0
    imgGet = np.array([[], []])  # 不能为空
    contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 检索外部轮廓
    Borderlist = []
    img = cv2.imread(input_img)
    imgCopy = img.copy()
    imgProcess = preProccessing(img)

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 800:  # 面积大于800像素为封闭图形
            cv2.drawContours(imgCopy, cnt, -1, (255, 0, 0), 3)
            peri = cv2.arcLength(cnt, True)  # 计算周长
            approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)  # 计算有多少个拐角
            x, y, w, h = cv2.boundingRect(approx)  # 得到外接矩形的大小
            a = (w + h) // 2
            dd = abs((w - h) // 2)  # 边框的差值
            imgGet = imgProcess[y:y + h, x:x + w]
            if w <= h:  # 得到一个正方形框，边界往外扩充20像素,黑色边框
                imgGet = cv2.copyMakeBorder(imgGet, 20, 20, 20 + dd, 20 + dd, cv2.BORDER_CONSTANT, value=[0, 0, 0])
                xx = x - dd - 10
                yy = y - 10
                ss = h + 20
                cv2.rectangle(imgCopy, (x - dd - 10, y - 10), (x + a + 10, y + h + 10), (0, 255, 0),
                              2)  # 看看框选的效果，在imgCopy中
            else:  # 边界往外扩充20像素值
                imgGet = cv2.copyMakeBorder(imgGet, 20 + dd, 20 + dd, 20, 20, cv2.BORDER_CONSTANT, value=[0, 0, 0])
                xx = x - 10
                yy = y - dd - 10
                ss = w + 20
                cv2.rectangle(imgCopy, (x - 10, y - dd - 10), (x + w + 10, y + a + 10), (0, 255, 0), 2)
            Temptuple = (imgGet, xx, yy, ss)  # 将图像及其坐标放在一个元组里面，然后再放进一个列表里面就可以访问了
            Borderlist.append(Temptuple)

    Borderlist.sort(key=lambda x: x[1])
    return Borderlist


def test(model_name, input_img="./test/image2.png", device="cuda"):
    if device == "cuda" and not torch.cuda.is_available():
        device = "cpu"
    print(f"Testing on {device}.")
    param_path = config.get_model_path(model_name)
    model = torch.load(param_path).to(device)
    model.eval()

    ans = ""
    index_to_class = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

    img = cv2.imread(input_img)
    imgProcessed = preProccessing(img)
    Borderlist = getContours(imgProcessed, input_img)

    train_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Grayscale(),
        transforms.Resize((28, 28)),
        transforms.ToTensor(),
    ])

    # if len(Borderlist) != 0:  # 不能为空
    if len(Borderlist) <= 1:
        image = Image.open(input_img)
        image = image.convert('L')
        image = image.resize((28, 28), Image.LANCZOS)
        image = np.array(image).reshape((1, 1, 28, 28)).astype('float32')
        image = torch.from_numpy(image).to(device)
        output = model(image)
        label = torch.argmax(output, 1).item()
        ans += str(label)
    else:
        for (imgRes, x, y, s) in Borderlist:
            img = train_transform(imgRes)
            img = torch.unsqueeze(img, dim=0).to(device)
            with torch.no_grad():
                pre = model(img)
                output = torch.squeeze(pre)
                predict = torch.softmax(output, dim=0)
                predict_cla = torch.argmax(predict).cpu().numpy()
                print(index_to_class[predict_cla], predict[predict_cla].cpu().numpy())
                result = index_to_class[predict_cla]
                ans += str(result)
    print("test answer = " + ans)
    return ans


if __name__ == "__main__":
    train(model_name="lenet", dataset=1)
